#!/usr/bin/env python3
import argparse
import grpc
import os
import sys
import struct
from time import sleep

# Import P4Runtime lib from parent utils dir
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "../utils/"))
import p4runtime_lib.bmv2
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections
import p4runtime_lib.helper

# 定义头部格式
ETH_HEADER_FORMAT = "!6s6sH"
CPU_HEADER_FORMAT = "!BB"
ETH_HEADER_LENGTH = struct.calcsize(ETH_HEADER_FORMAT)
CPU_HEADER_OFFSET = ETH_HEADER_LENGTH


def writeIpv4Rules(p4info_helper, sw, match_fields, action_params):
    """写入IPv4转发规则"""
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={"hdr.ipv4.dstAddr": match_fields},
        action_name="MyIngress.ipv4_forward",
        action_params=action_params,
    )
    sw.WriteTableEntry(table_entry)
    print(f"已为 {sw.name} 添加IPv4转发规则: {match_fields}")


def writeEcnRules(p4info_helper, sw, ecn_value, action_params):
    """写入ECN检查规则"""
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyEgress.check_ecn",
        match_fields={"hdr.ipv4.ecn": ecn_value},
        action_name="MyEgress.mark_ecn",
        action_params=action_params,
    )
    sw.WriteTableEntry(table_entry)
    print(f"已为 {sw.name} 添加ECN检查规则")


def writeCloneSession(p4info_helper, sw, session_id, ports):
    """设置克隆会话"""
    replicas = [{"egress_port": port, "instance": 1} for port in ports]
    clone_entry = p4info_helper.buildCloneSessionEntry(
        clone_session_id=session_id, replicas=replicas, packet_length_bytes=0
    )
    sw.WritePREEntry(clone_entry)
    print(f"已为 {sw.name} 设置克隆会话")


def monitorCongestion(connections):
    """监控网络拥塞"""
    print("开始监控网络拥塞...")
    congestion_stats = {conn.name: 0 for conn in connections}

    try:
        while True:
            sleep(2)
            print("\n" + "=" * 40)
            print("网络拥塞监控状态")
            print("=" * 40)

            for conn in connections:
                try:
                    for response in conn.stream_msg_resp:
                        if response.WhichOneof("update") == "packet":
                            packet = response.packet.payload

                            if len(packet) >= CPU_HEADER_OFFSET + 2:
                                cpu_header = packet[
                                    CPU_HEADER_OFFSET : CPU_HEADER_OFFSET + 2
                                ]
                                ecn, switch_id = struct.unpack(
                                    CPU_HEADER_FORMAT, cpu_header
                                )

                                print(f"{conn.name} 报告: ECN = {ecn}")

                                if ecn == 3:
                                    congestion_stats[conn.name] += 1
                                    print(f"🚨 🚨 🚨 警告: {conn.name} 检测到网络拥塞!")
                                    print(
                                        f"    拥塞次数: {congestion_stats[conn.name]}"
                                    )
                                    print(f"    建议调整流量控制策略")

                except grpc.RpcError:
                    pass

            print(f"\n拥塞统计: {congestion_stats}")
            print("." * 40)

    except KeyboardInterrupt:
        print("\n停止监控")


def main(p4info_file_path, bmv2_file_path):
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # 创建交换机连接
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name="s1",
            address="127.0.0.1:50051",
            device_id=0,
            proto_dump_file="logs/s1-p4runtime-requests.txt",
        )

        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name="s2",
            address="127.0.0.1:50052",
            device_id=1,
            proto_dump_file="logs/s2-p4runtime-requests.txt",
        )

        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name="s3",
            address="127.0.0.1:50053",
            device_id=2,
            proto_dump_file="logs/s3-p4runtime-requests.txt",
        )

        switches = [s1, s2, s3]

        # 获取用户配置
        print("ECN拥塞检测配置")
        print("-" * 25)
        try:
            threshold = int(input("请输入队列深度阈值 (默认20): ") or "20")
        except ValueError:
            threshold = 20

        print(f"设置ECN阈值为: {threshold}")

        # 建立控制器连接
        for switch in switches:
            switch.MasterArbitrationUpdate()

        # 安装P4程序
        for switch in switches:
            switch.SetForwardingPipelineConfig(
                p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path
            )
        print("P4程序安装完成")

        # 配置IPv4转发规则 - 修正版
        print("\n配置IPv4转发规则...")

        # s1的转发规则 - 连接到h1和h11
        writeIpv4Rules(
            p4info_helper,
            s1,
            ["10.0.1.1", 32],
            {"dstAddr": "08:00:00:00:01:01", "port": 2},
        )
        writeIpv4Rules(
            p4info_helper,
            s1,
            ["10.0.1.11", 32],
            {"dstAddr": "08:00:00:00:01:11", "port": 1},
        )
        writeIpv4Rules(
            p4info_helper,
            s1,
            ["10.0.2.0", 24],
            {"dstAddr": "08:00:00:00:02:00", "port": 3},
        )  # 到s2
        writeIpv4Rules(
            p4info_helper,
            s1,
            ["10.0.3.0", 24],
            {"dstAddr": "08:00:00:00:03:00", "port": 4},
        )  # 到s3

        # s2的转发规则 - 连接到h2和h22
        writeIpv4Rules(
            p4info_helper,
            s2,
            ["10.0.2.2", 32],
            {"dstAddr": "08:00:00:00:02:02", "port": 2},
        )
        writeIpv4Rules(
            p4info_helper,
            s2,
            ["10.0.2.22", 32],
            {"dstAddr": "08:00:00:00:02:22", "port": 1},
        )
        writeIpv4Rules(
            p4info_helper,
            s2,
            ["10.0.1.0", 24],
            {"dstAddr": "08:00:00:00:01:00", "port": 3},
        )  # 到s1
        writeIpv4Rules(
            p4info_helper,
            s2,
            ["10.0.3.0", 24],
            {"dstAddr": "08:00:00:00:03:00", "port": 4},
        )  # 到s3

        # s3的转发规则 - 连接到h3
        writeIpv4Rules(
            p4info_helper,
            s3,
            ["10.0.3.3", 32],
            {"dstAddr": "08:00:00:00:03:03", "port": 1},
        )
        writeIpv4Rules(
            p4info_helper,
            s3,
            ["10.0.1.0", 24],
            {"dstAddr": "08:00:00:00:01:00", "port": 2},
        )  # 到s1
        writeIpv4Rules(
            p4info_helper,
            s3,
            ["10.0.2.0", 24],
            {"dstAddr": "08:00:00:00:02:00", "port": 3},
        )  # 到s2

        # 配置ECN规则
        print("\n配置ECN检查规则...")
        for switch in switches:
            writeEcnRules(p4info_helper, switch, 1, {"ecn_threshold": threshold})
            writeEcnRules(p4info_helper, switch, 2, {"ecn_threshold": threshold})

        # 设置克隆
        print("\n配置克隆会话...")
        for switch in switches:
            writeCloneSession(p4info_helper, switch, 100, [252])

        print("\n配置完成! 开始监控...")
        monitorCongestion(switches)

    except KeyboardInterrupt:
        print("\n程序退出")
    except grpc.RpcError as e:
        printGrpcError(e)
    finally:
        ShutdownAllSwitchConnections()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="ECN拥塞监控控制器")
    parser.add_argument(
        "--p4info",
        help="p4info文件",
        type=str,
        action="store",
        required=False,
        default="./build/ecn.p4.p4info.txt",
    )
    parser.add_argument(
        "--bmv2-json",
        help="BMv2 JSON文件",
        type=str,
        action="store",
        required=False,
        default="./build/ecn.json",
    )
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print(f"\nP4Info文件不存在: {args.p4info}")
        print(
            "请先编译P4程序: p4c-bm2-ss --p4v 16 --p4runtime-files build/ecn.p4.p4info.txt -o build/ecn.json ecn.p4"
        )
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print(f"\nBMv2 JSON文件不存在: {args.bmv2_json}")
        print("请先编译P4程序")
        parser.exit(1)

    main(args.p4info, args.bmv2_json)
